"""
AI生成的代码
"""
import os
import shutil
import random

pic_suffix = (".png", ".jpg", ".jepg", ".bmp", ".tiff")


def split_data(
    image_folder: str,
    label_folder: str,
    train_ratio=0.7,
    test_ratio=0.2,
    output_dir=".",
):
    # 创建存储图片和标签的文件夹
    os.chdir(output_dir)
    os.makedirs("images", exist_ok=True)
    os.makedirs("labels", exist_ok=True)
    os.makedirs("images/train", exist_ok=True)
    os.makedirs("images/test", exist_ok=True)
    os.makedirs("images/val", exist_ok=True)
    os.makedirs("labels/train", exist_ok=True)
    os.makedirs("labels/test", exist_ok=True)
    os.makedirs("labels/val", exist_ok=True)

    # 获取图片文件名列表
    image_files = [f for f in os.listdir(image_folder) if any(f.endswith(s) for s in pic_suffix)]
    label_files = [f[:-4] + ".txt" for f in image_files]

    # assert len(image_files) == len(label_files), "图片和标签文件数量不匹配"
    # assert all(i[0][:-4] == i[1][:-4] for i in zip(image_files, label_files)), "图片与标签不对应"

    # 随机打乱图片文件名列表
    idx = list(range(len(image_files)))
    random.shuffle(idx)
    image_files = [image_files[i] for i in idx]
    label_files = [label_files[i] for i in idx]

    # 计算每个集合的大小
    total_size = len(image_files)
    train_size = int(total_size * train_ratio)
    test_size = int(total_size * test_ratio)

    # 将图片和标签文件分配到各个集合中
    train_images = image_files[:train_size]
    train_labels = label_files[:train_size]
    test_images = image_files[train_size : train_size + test_size]
    test_labels = label_files[train_size : train_size + test_size]
    val_images = image_files[train_size + test_size :]
    val_labels = label_files[train_size + test_size :]

    # 将图片和标签文件复制到相应的文件夹中
    for img, lbl in zip(train_images, train_labels):
        shutil.copy(os.path.join(image_folder, img), "images/train")
        shutil.copy(os.path.join(label_folder, lbl), "labels/train")
    for img, lbl in zip(test_images, test_labels):
        shutil.copy(os.path.join(image_folder, img), "images/test")
        shutil.copy(os.path.join(label_folder, lbl), "labels/test")
    for img, lbl in zip(val_images, val_labels):
        shutil.copy(os.path.join(image_folder, img), "images/val")
        shutil.copy(os.path.join(label_folder, lbl), "labels/val")


# 调用函数进行数据拆分
# split_data(r"担架", r"担架_yolo")
